package org.nerve.webmagic.downloader;

import org.nerve.webmagic.DriverBuilder;
import org.nerve.webmagic.bean.SpiderConfig;
import org.nerve.webmagic.utils.ScreenShotTool;
import org.openqa.selenium.By;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.WebElement;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import us.codecraft.webmagic.Page;
import us.codecraft.webmagic.Request;
import us.codecraft.webmagic.Site;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.downloader.Downloader;
import us.codecraft.webmagic.selector.PlainText;

import java.io.Closeable;
import java.io.IOException;

import static org.nerve.webmagic.bean.WebPage.SCREENSHOT;

/**
 *
 * org.nerve.webmagic.downloader
 * Created by zengxm on 2016/4/21 0021.
 */
public abstract class AbstractRemoteDriverDownloader<T extends WebDriver> implements Downloader, Closeable {
	protected Logger logger= LoggerFactory.getLogger(this.getClass());

	protected ThreadLocal<T> driverThreadLocal=new ThreadLocal<>();
	protected ThreadLocal<ScreenShotTool> screenThreadLocal=new ThreadLocal<>();

	protected final SpiderConfig config;
	public AbstractRemoteDriverDownloader(SpiderConfig config){
		this.config=config;
		driverPool = new DriverPool<T>(config, getDriverBuilder());
	}

	//WebDriver池
	protected DriverPool<T> driverPool;
	protected abstract DriverBuilder<T> getDriverBuilder();

	@Override
	public Page download(Request request, Task task) {
		logger.info("------------------------start download---------------------:threadName={}",Thread.currentThread().getName());
		final Site site=task.getSite();
		T webDriver=null;
		try {
//			webDriver = getDriver(site);
			webDriver=driverPool.getDriver();
		}  catch (InterruptedException e) {
			e.printStackTrace();
			logger.error("cannot get Driver:{}",e.getMessage());
			return null;
		}

		logger.info("[WebDriver={}]downloading page {}, thread={}",
				webDriver,
				request.getUrl(),
				Thread.currentThread().getName());

		//加载网页
		webDriver.get(request.getUrl());
		try {
			Thread.sleep(config.getSleepTime());
		} catch (InterruptedException e) {
			logger.error("try to sleep {} ms, but error on:{}", config.getSleepTime(), e.getMessage());
		}
		logger.info("[WebDriver={}]downloaded page {}", webDriver, request.getUrl());

		WebElement webElement = webDriver.findElement(By.xpath("/html"));
		String content = webElement.getAttribute("outerHTML");
		Page page = new Page();
		page.setRawText(content);
		page.setUrl(new PlainText(request.getUrl()));
		page.setRequest(request);

		//如果为true，则将截图赋值到 page 对象的 SCREENSHOT 上。
		if(config.isTakeScreenshot()){
			page.putField(SCREENSHOT, takeScreenShot(webDriver));
		}

		driverPool.returnToPool(webDriver);
		logger.info("[WebDriver={}] load data of page, next to processor {}", webDriver, request.getUrl());
		return page;
	}

	/**
	 * 进行网页截图，不同的Driver可以重写此方法
	 * @param driver
	 * @return
	 */
	protected String takeScreenShot(final T driver){
		ScreenShotTool tool=screenThreadLocal.get();
		if(tool==null){
			tool=new ScreenShotTool();
			screenThreadLocal.set(tool);
		}

		return tool.taskScreenshot(driver);
	}

	@Override
	public void setThread(int i) {
		config.setThreadSize(i);
	}

	@Override
	public void close() throws IOException {
		driverPool.close();
	}
}
